import os
import cv2
import json
import datetime
import numpy as np
import time
import threading
import shutil
import logging
import logging.handlers
from .rerun_visualizer import RerunLogger
from queue import Queue, Empty
from threading import Thread, Event, Lock
from enum import Enum

# 配置日志系统
def setup_logger(name, log_file=None, level=logging.INFO):
    """设置日志记录器，支持同时输出到控制台和文件"""
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.propagate = False  # 避免日志传播到根记录器
    
    # 清除现有处理器
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)
    
    # 创建控制台处理器 - 仅显示 WARNING 及以上级别
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.WARNING)
    console_format = logging.Formatter('%(levelname)s - %(name)s - %(message)s')
    console_handler.setFormatter(console_format)
    logger.addHandler(console_handler)
    
    # 如果指定了日志文件，创建文件处理器
    if log_file:
        # 确保日志目录存在
        log_dir = os.path.dirname(log_file)
        if log_dir and not os.path.exists(log_dir):
            os.makedirs(log_dir)
            
        # 创建按大小轮转的文件处理器
        file_handler = logging.handlers.RotatingFileHandler(
            log_file, maxBytes=10*1024*1024, backupCount=5, encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        file_format = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(file_format)
        logger.addHandler(file_handler)
    
    return logger

# 创建日志记录器
logger = setup_logger('EpisodeWriter', log_file='logs/episode_writer.log')

# 定义 EpisodeWriter 的状态枚举
class EpisodeState(Enum):
    IDLE = 0       # 空闲状态，可以创建新的 episode
    RECORDING = 1  # 正在记录数据
    SAVING = 2     # 正在保存数据
    ERROR = 3      # 错误状态

class EpisodeWriter():
    def __init__(self, task_dir, frequency=30, image_size=[640, 480], rerun_log=True, save_timeout=60):
        """
        初始化 EpisodeWriter
        
        Args:
            task_dir: 数据保存目录
            frequency: 数据记录频率
            image_size: 图像尺寸 [width, height]
            rerun_log: 是否启用 Rerun 可视化
            save_timeout: 保存超时时间（秒）
        """
        logger.info("EpisodeWriter 初始化中...")
        self.task_dir = task_dir
        self.frequency = frequency
        self.image_size = image_size
        self.save_timeout = save_timeout

        # Rerun 可视化初始化
        self.rerun_log = rerun_log
        if self.rerun_log:
            logger.info("RerunLogger 初始化中...")
            self.rerun_logger = RerunLogger(prefix="online/", IdxRangeBoundary=60, memory_limit="300MB")
            logger.info("RerunLogger 初始化完成")
        
        # 数据存储相关初始化
        self.data = {}
        self.episode_data = []
        self.item_id = -1
        self.episode_id = -1
        
        # 检查并创建任务目录，确定 episode_id
        if os.path.exists(self.task_dir):
            episode_dirs = [episode_dir for episode_dir in os.listdir(self.task_dir) if 'episode_' in episode_dir]
            episode_last = sorted(episode_dirs)[-1] if len(episode_dirs) > 0 else None
            self.episode_id = 0 if episode_last is None else int(episode_last.split('_')[-1])
            logger.info(f"任务目录已存在，当前 episode_id: {self.episode_id}")
        else:
            os.makedirs(self.task_dir)
            logger.info("任务目录不存在，已创建新目录")
        
        # 初始化数据信息
        self.data_info()
        self.text_desc()

        # 状态管理
        self._state = EpisodeState.IDLE
        self._state_lock = Lock()  # 状态锁，保证状态修改的原子性
        self._save_event = Event()  # 保存完成事件
        
        # 统计信息
        self.stats = {
            'total_episodes': 0,           # 总 episode 数
            'total_items': 0,              # 总数据项数
            'total_images': 0,             # 总图像数
            'total_save_time': 0,          # 总保存时间(秒)
            'avg_save_time': 0,            # 平均保存时间(秒)
            'last_save_time': 0,           # 最近一次保存时间(秒)
            'max_save_time': 0,            # 最长保存时间(秒)
            'save_count': 0,               # 保存次数
            'error_count': 0,              # 错误次数
            'start_time': time.time(),     # 启动时间
        }
        
        # 队列和线程初始化
        self.item_data_queue = Queue(maxsize=100)  # 数据项队列
        self.save_queue = Queue(maxsize=5)         # 保存任务队列
        self.stop_worker = False
        
        # 启动工作线程
        self.worker_thread = Thread(target=self.process_queue, name="DataProcessThread")
        self.save_thread = Thread(target=self.process_save, name="SaveProcessThread")
        self.worker_thread.daemon = True  # 设为守护线程，主程序退出时自动结束
        self.save_thread.daemon = True
        self.worker_thread.start()
        self.save_thread.start()

        logger.info("EpisodeWriter 初始化完成")
        
    def get_stats(self):
        """获取当前统计信息"""
        # 计算运行时间
        runtime = time.time() - self.stats['start_time']
        
        # 构建完整的统计信息
        stats = self.stats.copy()
        stats.update({
            'runtime': runtime,
            'current_state': self.state.name,
            'data_queue_size': self.item_data_queue.qsize(),
            'save_queue_size': self.save_queue.qsize(),
            'current_episode_id': self.episode_id,
            'current_item_id': self.item_id,
        })
        
        return stats
        
    def print_stats(self):
        """打印当前统计信息"""
        stats = self.get_stats()
        
        print("\n===== EpisodeWriter 统计信息 =====")
        print(f"\u8fd0行时间: {stats['runtime']:.2f} 秒")
        print(f"\u5f53前状态: {stats['current_state']}")
        print(f"\u5f53前 Episode ID: {stats['current_episode_id']}")
        print(f"\u5f53前 Item ID: {stats['current_item_id']}")
        print(f"\u603b Episode 数: {stats['total_episodes']}")
        print(f"\u603b数据项数: {stats['total_items']}")
        print(f"\u603b图像数: {stats['total_images']}")
        print(f"\u6570据队列大小: {stats['data_queue_size']}")
        print(f"\u4fdd存队列大小: {stats['save_queue_size']}")
        print(f"\u4fdd存次数: {stats['save_count']}")
        print(f"\u5e73均保存时间: {stats['avg_save_time']:.2f} 秒")
        print(f"\u6700长保存时间: {stats['max_save_time']:.2f} 秒")
        print(f"\u6700近保存时间: {stats['last_save_time']:.2f} 秒")
        print(f"\u9519误次数: {stats['error_count']}")
        print("==============================\n")
        
    @property
    def state(self):
        """获取当前状态"""
        with self._state_lock:
            return self._state
            
    @property
    def is_available(self):
        """检查是否可以创建新的 episode"""
        return self.state == EpisodeState.IDLE
        
    def _set_state(self, new_state):
        """安全地设置状态"""
        with self._state_lock:
            old_state = self._state
            self._state = new_state
            logger.debug(f"状态变更: {old_state} -> {new_state}")
            return old_state

    def data_info(self, version='1.0.0', date=None, author=None):
        self.info = {
                "version": "1.0.0" if version is None else version, 
                "date": datetime.date.today().strftime('%Y-%m-%d') if date is None else date,
                "author": "unitree" if author is None else author,
                "image": {"width":self.image_size[0], "height":self.image_size[1], "fps":self.frequency},
                "depth": {"width":self.image_size[0], "height":self.image_size[1], "fps":self.frequency},
                "audio": {"sample_rate": 16000, "channels": 1, "format":"PCM", "bits":16},    # PCM_S16
                "joint_names":{
                    "left_arm":   ['kLeftShoulderPitch' ,'kLeftShoulderRoll', 'kLeftShoulderYaw', 'kLeftElbow', 'kLeftWristRoll', 'kLeftWristPitch', 'kLeftWristyaw'],
                    "left_hand":  [],
                    "right_arm":  [],
                    "right_hand": [],
                    "body":       [],
                },

                "tactile_names": {
                    "left_hand": [],
                    "right_hand": [],
                }, 
            }
    def text_desc(self):
        self.text = {
            "goal": "Pick up the red cup on the table.",
            "desc": "Pick up the cup from the table and place it in another position. The operation should be smooth and the water in the cup should not spill out",
            "steps":"step1: searching for cups. step2: go to the target location. step3: pick up the cup",
        }

 
    def create_episode(self, timeout=5.0):
        """
        创建新的数据集片段（episode）
        
        Args:
            timeout: 等待可用状态的超时时间（秒）
            
        Returns:
            bool: 创建成功返回 True，失败返回 False
        """
        # 等待可用状态，最多等待 timeout 秒
        start_time = time.time()
        while not self.is_available:
            if time.time() - start_time > timeout:
                logger.warning(f"等待可用状态超时（{timeout}秒），无法创建新的 episode")
                return False
            time.sleep(0.1)
        
        # 切换到记录状态
        if self._set_state(EpisodeState.RECORDING) != EpisodeState.IDLE:
            logger.error(f"无法从当前状态切换到记录状态: {self.state}")
            return False
        
        # 初始化数据结构
        self.item_id = -1
        self.episode_data = []
        self.episode_id += 1
        
        # 创建目录结构
        self.episode_dir = os.path.join(self.task_dir, f"episode_{str(self.episode_id).zfill(4)}")
        try:
            # 创建主目录
            os.makedirs(self.episode_dir, exist_ok=True)
            
            # 创建各类数据子目录
            self.color_dir = os.path.join(self.episode_dir, 'colors')
            self.depth_dir = os.path.join(self.episode_dir, 'depths')
            self.audio_dir = os.path.join(self.episode_dir, 'audios')
            os.makedirs(self.color_dir, exist_ok=True)
            os.makedirs(self.depth_dir, exist_ok=True)
            os.makedirs(self.audio_dir, exist_ok=True)
            
            # 设置 JSON 文件路径
            self.json_path = os.path.join(self.episode_dir, 'data.json')
            
            logger.info(f"成功创建新的 episode: {self.episode_id}")
            return True
            
        except Exception as e:
            logger.error(f"创建 episode 目录结构失败: {e}")
            self._set_state(EpisodeState.ERROR)
            return False

    def add_item(self, colors, depths=None, states=None, actions=None, tactiles=None, audios=None):
        """
        添加一个数据项到当前 episode
        
        Args:
            colors: 彩色图像字典 {color_key: image_array}
            depths: 深度图字典 {depth_key: depth_array}
            states: 状态数据字典
            actions: 动作数据字典
            tactiles: 触觉数据字典
            audios: 音频数据字典
            
        Returns:
            bool: 添加成功返回 True，失败返回 False
        """
        # 检查当前是否处于记录状态
        if self.state != EpisodeState.RECORDING:
            logger.warning(f"当前状态不是记录状态，无法添加数据项: {self.state}")
            return False
            
        # 递增数据项 ID
        self.item_id += 1
        
        # 创建数据项字典
        item_data = {
            'idx': self.item_id,
            'colors': colors,
            'depths': depths,
            'states': states,
            'actions': actions,
            'tactiles': tactiles,
            'audios': audios,
            'timestamp': time.time()
        }
        
        # 将数据项加入队列
        try:
            self.item_data_queue.put(item_data, timeout=0.5)  # 设置超时，避免队列满时阻塞
            logger.debug(f"数据项已加入队列: {self.item_id}")
            return True
        except Exception as e:
            logger.error(f"添加数据项到队列失败: {e}")
            return False

    def process_queue(self):
        """数据处理线程的主函数。处理队列中的数据项并保存到磁盘"""
        logger.info("数据处理线程启动")
        
        while not self.stop_worker:
            try:
                # 从队列中获取数据项
                try:
                    item_data = self.item_data_queue.get(timeout=1)
                    
                    # 处理数据项
                    try:
                        self._process_item_data(item_data)
                        logger.debug(f"数据项处理成功: {item_data['idx']}")
                    except Exception as e:
                        logger.error(f"处理数据项失败 (idx={item_data['idx']}): {e}")
                    finally:
                        # 标记任务完成，无论成功与否
                        self.item_data_queue.task_done()
                        
                except Empty:
                    # 队列为空，等待新数据
                    pass
                    
            except Exception as e:
                logger.error(f"数据处理线程异常: {e}")
                time.sleep(1)  # 避免循环过快消耗 CPU
                
        logger.info("数据处理线程结束")

    def _process_item_data(self, item_data):
        """
        处理单个数据项，将图像、深度图和音频保存到磁盘
        
        Args:
            item_data: 数据项字典
        """
        idx = item_data['idx']
        colors = item_data.get('colors', {})
        depths = item_data.get('depths', {})
        audios = item_data.get('audios', {})

        # 保存彩色图像
        if colors:
            for idx_color, (color_key, color) in enumerate(colors.items()):
                color_name = f'{str(idx).zfill(6)}_{color_key}.jpg'
                color_path = os.path.join(self.color_dir, color_name)
                try:
                    if not cv2.imwrite(color_path, color):
                        logger.error(f"保存彩色图像失败: {color_path}")
                    item_data['colors'][color_key] = os.path.join('colors', color_name)
                except Exception as e:
                    logger.error(f"保存彩色图像异常: {color_path}, {e}")

        # 保存深度图
        if depths:
            for idx_depth, (depth_key, depth) in enumerate(depths.items()):
                depth_name = f'{str(idx).zfill(6)}_{depth_key}.jpg'
                depth_path = os.path.join(self.depth_dir, depth_name)
                try:
                    if not cv2.imwrite(depth_path, depth):
                        logger.error(f"保存深度图失败: {depth_path}")
                    item_data['depths'][depth_key] = os.path.join('depths', depth_name)
                except Exception as e:
                    logger.error(f"保存深度图异常: {depth_path}, {e}")

        # 保存音频数据
        if audios:
            for mic, audio in audios.items():
                audio_name = f'audio_{str(idx).zfill(6)}_{mic}.npy'
                audio_path = os.path.join(self.audio_dir, audio_name)
                try:
                    np.save(audio_path, audio.astype(np.int16))
                    item_data['audios'][mic] = os.path.join('audios', audio_name)
                except Exception as e:
                    logger.error(f"保存音频数据异常: {audio_path}, {e}")

        # 更新 episode 数据
        self.episode_data.append(item_data)

        # 如果启用了 Rerun 日志，记录数据
        if self.rerun_log and hasattr(self, 'rerun_logger'):
            try:
                self.rerun_logger.log_item_data(item_data)
                logger.debug(f"数据项已记录到 Rerun: {idx}")
            except Exception as e:
                logger.error(f"记录数据到 Rerun 失败: {e}")
                
    def process_save(self):
        """保存线程的主函数。处理保存队列中的保存请求"""
        logger.info("保存线程启动")
        
        while not self.stop_worker:
            try:
                # 检查是否有保存请求
                try:
                    # 从保存队列中获取保存任务
                    save_task = self.save_queue.get(timeout=1)
                    
                    # 执行保存操作
                    try:
                        self._save_episode(save_task)
                        logger.info(f"成功保存 episode {save_task['episode_id']}")
                    except Exception as e:
                        logger.error(f"保存 episode 失败: {e}")
                    finally:
                        # 标记保存任务完成
                        self.save_queue.task_done()
                        
                except Empty:
                    # 保存队列为空，等待新的保存请求
                    pass
                    
            except Exception as e:
                logger.error(f"保存线程异常: {e}")
                time.sleep(1)  # 避免循环过快消耗 CPU
                
        logger.info("保存线程结束")

    def save_episode(self, wait_timeout=None):
        """
        触发保存操作，将保存请求加入队列
        
        Args:
            wait_timeout: 等待保存完成的超时时间（秒），如果为 None 则不等待
            
        Returns:
            bool: 保存请求发送成功返回 True，失败返回 False
        """
        # 检查当前是否处于记录状态
        if self.state != EpisodeState.RECORDING:
            logger.warning(f"当前状态不是记录状态，无法触发保存: {self.state}")
            return False
            
        # 切换到保存状态
        self._set_state(EpisodeState.SAVING)
        
        # 清除之前的保存完成事件
        self._save_event.clear()
        
        # 准备保存任务
        save_task = {
            'episode_id': self.episode_id,
            'episode_dir': self.episode_dir,
            'json_path': self.json_path,
            'episode_data': self.episode_data.copy(),  # 复制数据，避免竞争条件
            'info': self.info,
            'text': self.text
        }
        
        # 将保存任务加入队列
        try:
            self.save_queue.put(save_task, timeout=0.5)
            logger.info(f"保存请求已加入队列: episode {self.episode_id}")
            
            # 如果指定了超时时间，等待保存完成
            if wait_timeout is not None:
                if self._save_event.wait(timeout=wait_timeout):
                    logger.info(f"保存完成，等待时间: {wait_timeout} 秒")
                else:
                    logger.warning(f"等待保存完成超时: {wait_timeout} 秒")
            
            return True
            
        except Exception as e:
            logger.error(f"发送保存请求失败: {e}")
            # 恢复到记录状态
            self._set_state(EpisodeState.RECORDING)
            return False

    def _save_episode(self, save_task):
        """
        实际执行保存操作，使用临时文件和原子重命名确保数据一致性
        
        Args:
            save_task: 保存任务字典，包含要保存的所有数据
        """
        episode_id = save_task['episode_id']
        json_path = save_task['json_path']
        episode_data = save_task['episode_data']
        
        # 开始计时
        save_start_time = time.time()
        
        # 准备保存数据
        save_data = {
            'info': save_task['info'],
            'text': save_task['text'],
            'data': episode_data
        }
        
        # 使用临时文件
        temp_path = f"{json_path}.tmp"
        
        try:
            # 先写入临时文件
            with open(temp_path, 'w', encoding='utf-8') as jsonf:
                jsonf.write(json.dumps(save_data, indent=4, ensure_ascii=False))
            
            # 使用原子重命名操作替换最终文件
            os.replace(temp_path, json_path)
            
            # 计算保存时间
            save_time = time.time() - save_start_time
            
            # 更新统计信息
            self.stats['total_episodes'] += 1
            self.stats['total_items'] += len(episode_data)
            self.stats['save_count'] += 1
            self.stats['last_save_time'] = save_time
            self.stats['total_save_time'] += save_time
            self.stats['avg_save_time'] = self.stats['total_save_time'] / self.stats['save_count']
            self.stats['max_save_time'] = max(self.stats['max_save_time'], save_time)
            
            # 统计图像数量
            image_count = 0
            for item in episode_data:
                if item.get('colors'):
                    image_count += len(item.get('colors', {}))
                if item.get('depths'):
                    image_count += len(item.get('depths', {}))
            self.stats['total_images'] += image_count
            
            logger.info(f"成功保存数据到 {json_path}, 用时: {save_time:.2f}秒, 数据项: {len(episode_data)}, 图像: {image_count}")
            
        except Exception as e:
            # 更新错误统计
            self.stats['error_count'] += 1
            
            logger.error(f"保存数据到 {json_path} 失败: {e}")
            # 尝试清理临时文件
            try:
                if os.path.exists(temp_path):
                    os.remove(temp_path)
            except:
                pass
            raise  # 重新抛出异常以便上层处理
            
        finally:
            # 无论成功与否，都将状态设置回空闲状态
            self._set_state(EpisodeState.IDLE)
            # 触发保存完成事件
            self._save_event.set()

    def close(self):
        """
        停止工作线程并确保所有任务完成
        """
        logger.info("关闭 EpisodeWriter...")
        
        # 等待数据项队列处理完成
        try:
            self.item_data_queue.join()
            logger.info("数据项队列处理完成")
        except Exception as e:
            logger.error(f"等待数据项队列完成异常: {e}")
        
        # 如果当前处于记录状态，触发保存
        if self.state == EpisodeState.RECORDING:
            logger.info("检测到未保存的数据，触发保存")
            self.save_episode(wait_timeout=10.0)  # 等待最多 10 秒
        
        # 等待保存队列处理完成
        try:
            self.save_queue.join()
            logger.info("保存队列处理完成")
        except Exception as e:
            logger.error(f"等待保存队列完成异常: {e}")
        
        # 停止工作线程
        self.stop_worker = True
        
        # 等待线程结束
        try:
            self.worker_thread.join(timeout=5.0)
            self.save_thread.join(timeout=5.0)
            logger.info("线程已安全结束")
        except Exception as e:
            logger.error(f"等待线程结束异常: {e}")
            
        logger.info("EpisodeWriter 已安全关闭")
        
    # 兼容性方法，为了支持旧版代码
    @property
    def need_save(self):
        return self.state == EpisodeState.SAVING
        
    @need_save.setter
    def need_save(self, value):
        # 兼容旧版代码的 need_save = True/False 设置
        if value and self.state == EpisodeState.RECORDING:
            # 触发保存
            self.save_episode()
        elif not value and self.state == EpisodeState.SAVING:
            # 重置保存状态
            self._set_state(EpisodeState.IDLE)